import sys
sys.path.append('./')
import torch
import _pickle as cPickle
import numpy as np
import os
import time
from argparse import ArgumentParser
from mmdet.apis import inference_detector, init_detector, show_result_pyplot

image_root = '/home/robot/rs_imgs'
cat_id = 41
# cat       coco_id     nocs_id
# cup       41          5
# bowl      45          1
# bottle    39          0
coco_list = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
            'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
            'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
            'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
            'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
            'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
            'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
            'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
            'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
            'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
id_map = {
    41: 5,
    45: 1,
    39: 0
}

def get_args():
    parser = ArgumentParser()
    parser.add_argument('--img_id', type=int, default=0,help='Image file')
    # parser.add_argument('--config', default='mask/configs/queryinst/queryinst.py',help='Config file')
    # parser.add_argument('--checkpoint', default='mask/queryinst_weights.pth',help='Checkpoint file')
    parser.add_argument('--config', default='vision_mask/configs/solo/decoupled_solo_light_r50_fpn_3x_coco.py',help='Config file')
    parser.add_argument('--checkpoint', default='vision_mask/decoupled_solo_weights.pth',help='Checkpoint file')
    parser.add_argument('--device', default='cuda:0', help='Device used for inference')
    parser.add_argument('--score-thr', type=float, default=0.3, help='bbox score threshold')
    parser.add_argument('--show_mask', action='store_true', default=False)
    parser.add_argument('--save_mask', type=bool, default=False)
    args = parser.parse_args()
    return args

def get_bbox_from_mask(mask):
    # mask: bool array, shape: (H, W)
    # return: (x1, y1, x2, y2)
    mask_ids = np.argwhere(mask == True)
    left_top = mask_ids.min(axis=0)
    right_bottom = mask_ids.max(axis=0)
    return np.array([left_top[1], left_top[0], right_bottom[1], right_bottom[0]])

def get_mask(img, model, args):
    if model is None:
        model = init_detector(args.config, args.checkpoint, device=args.device)
    result_ins = {}
    result = inference_detector(model, img)
    # show the results
    if args.show_mask:
        show_result_pyplot(model, img, result, score_thr=args.score_thr)
    # Here, we only consider the scene with one object.
    max_score = 0.
    for id in id_map.keys():
        if result[0][id].shape[0] > 0:
            score = result[0][id][:, 4].max()
            if score > max_score:
                max_score = score
                cat_id = id
    result_ins['class_name'] = coco_list[cat_id]
    result_ins['class_ids'] = id_map[cat_id]
    result_ins['scores'] = result[0][cat_id][:, 4]
    best_idx = result[0][cat_id][:, 4].argmax()
    result_ins['scores'] = result_ins['scores'][best_idx]
    result_ins['masks'] = np.array(result[1][cat_id])[best_idx]
    # result_ins['rois'] = result[0][cat_id][best_idx, :4]
    # get bbox from mask:
    result_ins['rois'] = get_bbox_from_mask(result_ins['masks'])
    return result_ins

# yjh,  2021/12/21
def get_mask_multi_obj(img, model, args):
    if model is None:
        model = init_detector(args.config, args.checkpoint, device=args.device)
    result = inference_detector(model, img)
    # show the results
    if args.show_mask:
        show_result_pyplot(model, img, result, score_thr=args.score_thr)
    # Here, we consider the scene with multi-objects of different categories.
    score_thres = 0.2
    results_ins = {'class_name':[], 'class_ids':[], 'scores':[], 'masks': [], 'rois':[]}
    for id in id_map.keys():
        if result[0][id].shape[0] > 0:
            score = result[0][id][:, 4].max()
            if score > score_thres:
                cat_id = id
                results_ins['class_name'].append(coco_list[cat_id])
                results_ins['class_ids'].append(id_map[cat_id])
                best_idx = result[0][cat_id][:, 4].argmax()
                results_ins['scores'].append(score)
                mask = np.array(result[1][cat_id])[best_idx]
                results_ins['masks'].append(mask)
                results_ins['rois'].append(get_bbox_from_mask(mask))
    for k, v in results_ins.items():
        results_ins[k] = np.array(v)
    return results_ins

def main():
    result_ins = {}
    args = get_args()
    model = init_detector(args.config, args.checkpoint, device=args.device)
    start_time_all = time.time()
    for i in range(args.img_id, args.img_id + 17):
        img_pth = os.path.join(image_root, '%04d' % i + '_color.png')
        t_start = time.time()
        result_ins = get_mask(img_pth, model, args)
        inference_time = time.time() - t_start
        print("%04d has done" % i, 'time: {}, cat_name: {}'.format(inference_time, result_ins['class_name']))
        
        if args.save_mask:
            with open(os.path.join(image_root, '%04d' % i + '_mask.pkl'), 'wb') as f:
                cPickle.dump(result_ins, f)
    print("Total time: {}, Average time: {}".format(time.time() - start_time_all, (time.time() - start_time_all) / (i+1)))


if __name__ == '__main__':
    main()
